package balance

import (
	"context"
	"fmt"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"strconv"
	"sync"
	"time"
)

const name = "custom_wrr"

func init() {
	balancer.Register(base.NewBalancerBuilder(name, &PickBuilder{}, base.Config{HealthCheck: true}))
}

type PickBuilder struct {
	oldConns []*conn
}

func (p *PickBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
	cs := info.ReadySCs
	conns := make([]*conn, 0, len(cs))
	for cc, ci := range cs {
		cn := &conn{cc: cc, ci: ci}
		metadata, ok := ci.Address.Metadata.(map[string]string)
		if ok {
			value := metadata["weight"]
			val, err := strconv.Atoi(value)
			if err == nil {
				cn.weight = val
			}
		}
		if cn.weight == 0 {
			//设置初始化
			cn.weight = 10
		}
		conns = append(conns, cn)
	}
	//if p.oldConns == nil{
	//	p.oldConns = conns
	//	return &Pick{conns: conns}
	//}else{
	//
	//}
	return &Pick{conns: conns}
}

type Pick struct {
	conns []*conn
	mutex sync.Mutex
}

type conn struct {
	cc            balancer.SubConn
	ci            base.SubConnInfo
	weight        int
	currentWeight int
	mutex         sync.Mutex

	efficientWeight int

	// 可以用来标记不可用（比如说熔断了）
	available bool
	// 记录标记为不可用的标记时间
	expireTime time.Time
}

func (p *Pick) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	var total int
	if len(p.conns) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}
	p.mutex.Lock()
	maxCC := p.conns[0]
	for _, conn := range p.conns {
		total = total + conn.weight
		conn.currentWeight = conn.weight + conn.currentWeight
		if maxCC == nil || maxCC.currentWeight < conn.currentWeight {
			maxCC = conn
		}
	}
	maxCC.currentWeight = maxCC.currentWeight - total
	fmt.Println(maxCC.ci.Address.Addr)
	p.mutex.Unlock()
	return balancer.PickResult{SubConn: maxCC.cc,
		Done: func(info balancer.DoneInfo) {
			err := info.Err
			if err == nil {
				return
			}
			switch err {
			case context.Canceled:
				return
			case context.DeadlineExceeded:

			default:
				//st, ok := status.FromError(err)

			}

		}}, nil
}

func (w *Pick) PickV1(info balancer.PickInfo) (balancer.PickResult, error) {
	if len(w.conns) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}
	var totalWeight int
	var res *conn
	//w.mutex.Lock()
	//defer w.mutex.Unlock()
	for _, c := range w.conns {
		c.mutex.Lock()
		totalWeight = totalWeight + c.efficientWeight
		c.currentWeight = c.currentWeight + c.efficientWeight
		if res == nil || res.currentWeight < c.currentWeight {
			res = c
		}
		c.mutex.Unlock()
	}
	res.mutex.Lock()
	res.currentWeight = res.currentWeight - totalWeight
	res.mutex.Unlock()
	return balancer.PickResult{
		SubConn: res.cc,
		Done: func(info balancer.DoneInfo) {
			res.mutex.Lock()
			defer res.mutex.Unlock()

			if info.Err != nil && res.efficientWeight == 1 {
				return
			}
			// MaxUint32 可以替换为你认为的最大值。
			// 例如说你预期节点的权重是在 100 - 200 之间
			// 那么你可以设置经过动态调整之后的权重不会超过 500。
			if info.Err == nil && res.efficientWeight >= 500 {
				return
			}
			if info.Err != nil {
				res.efficientWeight--
			} else {
				res.efficientWeight++
			}
		},
	}, nil
}

func (w *Pick) PickV2(info balancer.PickInfo) (balancer.PickResult, error) {
	if len(w.conns) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}
	var totalWeight int
	var res *conn
	//w.mutex.Lock()
	//defer w.mutex.Unlock()
	for _, c := range w.conns {
		c.mutex.Lock()
		totalWeight = totalWeight + c.efficientWeight
		c.currentWeight = c.currentWeight + c.efficientWeight
		if res == nil || res.currentWeight < c.currentWeight {
			res = c
		}
		c.mutex.Unlock()
	}
	res.mutex.Lock()
	res.currentWeight = res.currentWeight - totalWeight
	res.mutex.Unlock()
	return balancer.PickResult{
		SubConn: res.cc,
		Done: func(info balancer.DoneInfo) {
			res.mutex.Lock()
			defer res.mutex.Unlock()

			if info.Err != nil && res.efficientWeight == 0 {
				return
			}
			// MaxUint32 可以替换为你认为的最大值。
			// 例如说你预期节点的权重是在 100 - 200 之间
			// 那么你可以设置经过动态调整之后的权重不会超过 500。
			switch info.Err {
			case nil:
				if res.efficientWeight == 500 {
					return
				}
				res.efficientWeight++
			case context.DeadlineExceeded: //超时
				// 超时可以考虑动态调整。
				// 比如说第一次超时是降低 1，第二次连续超时是降低 2
				res.efficientWeight = res.efficientWeight - 10
			default:
				code := status.Code(info.Err)
				switch code {
				// 直接降低到 1，我们可以预期接下来几乎不会选中它。
				// 但是本身没有降低到 0，所以它又存在被选中的机会，
				// 那么后续会慢慢恢复过来
				case codes.Unavailable:
					//服务不可用  直接讲到最低，没有降到0，所以它存在被选中的机会，会慢慢恢复过来
					res.efficientWeight = 1

					//这里也可以采用直接编为不可用，然后过段时间在恢复过来，记录一下不可用时间。
					//res.available=false
					//res.expireTime=time.Now()

				case codes.ResourceExhausted:
					//限流 直接减半，与tcp流量控制一样，降半恢复。
					res.efficientWeight = res.efficientWeight / 2
				case codes.Aborted:
					//降级 也是直接减半
					res.efficientWeight = res.efficientWeight / 2
				default:
					//其他错误
					if res.efficientWeight == 1 {
						return
					}
					res.efficientWeight--
				}
			}
		},
	}, nil
}
